package com.yangtao.prototype.common.util;

import com.yangtao.framework.util.PropertyUtil;
import org.apache.log4j.Logger;
import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.DatabaseSequenceFilter;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.FilteredDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.filter.ITableFilter;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;

import java.io.File;
import java.io.FileOutputStream;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Map;
import java.util.Properties;
import java.util.Set;

/**
 * 将xml或者excel中的数据同步到数据库中
 * 同时也可以执行sql语句
 *
 * @author fudongming
 * @date: 11-6-3 - 上午10:21
 */
public class DbUnitUtil {
	private String driverName;
	private String jdbcUrl;
	private String username;
	private String password;
	/**
	 * 原生的连接
	 */
	private Connection connection;
	private static final Logger logger = Logger.getLogger(DbUnitUtil.class);

	/**
	 * 默认的初始化连接参数的方法，如果不用这个方法，
	 * 那么必须指定连接参数
	 */
	public void initialize() {
		Properties jdbc = PropertyUtil.getProperties("/jdbc.properties");
		this.driverName = jdbc.getProperty("jdbc.driverClassName");
		this.jdbcUrl = jdbc.getProperty("jdbc.url");
		this.username = jdbc.getProperty("jdbc.username");
		this.password = jdbc.getProperty("jdbc.password");
	}

	/**
	 * 获取包装过的连接
	 *
	 * @return
	 */
	public IDatabaseConnection createConnect() {
		if (logger.isInfoEnabled()) {
			logger.info("create dbunit connection: " + jdbcUrl);
		}
		try {
			Connection jdbcConnection = null;
			if (connection == null) {
				Class.forName(driverName);
				jdbcConnection = DriverManager.getConnection(
						jdbcUrl, username, password);
			} else {
				jdbcConnection = connection;
			}
			IDatabaseConnection connection = new DatabaseConnection(jdbcConnection);
			return connection;
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		} catch (SQLException e) {
			e.printStackTrace();
		} catch (DatabaseUnitException e) {
			e.printStackTrace();
		}
		if (logger.isInfoEnabled()) {
			logger.info("connection create success");
		}
		return null;
	}

	/**
	 * 数据过滤器
	 *
	 * @return 数据过滤器
	 * @throws Exception
	 */
	private ITableFilter createFilter() throws Exception {
		ITableFilter filter = new DatabaseSequenceFilter(createConnect());
		return filter;
	}

	public ITable query(String tableName, String querySQL) throws Exception {
		IDatabaseConnection connection = null;
		connection = createConnect();
		if (logger.isInfoEnabled()) {
			logger.info("execute query: " + querySQL);
		}
		ITable data = connection.createQueryTable(tableName, querySQL);
		return data;
	}

	public void executeSql(String sql, Object... parameters) throws Exception {
		IDatabaseConnection connection = null;
		PreparedStatement pst = null;
		try {
			connection = createConnect();
			if (logger.isInfoEnabled()) {
				logger.info("execute query: " + sql);
			}
			pst = connection.getConnection().prepareStatement(sql);
			if (parameters != null && parameters.length > 0) {
				for (int i = 1; i <= parameters.length; i++) {
					pst.setObject(i, parameters[i - 1]);
				}
			}
			pst.executeUpdate();
		} finally {
			if (pst != null) {
				pst.close();
			}
		}
	}

	public void exportTableDataToXML(Map<String, File> tableToXML) throws Exception {
		IDatabaseConnection connection = null;

		try {
			connection = createConnect();

			Set<String> tableSet = tableToXML.keySet();
			for (String tableName : tableSet) {
				QueryDataSet partialDataSet = new QueryDataSet(connection);
				partialDataSet.addTable(tableName);
				FlatXmlDataSet.write(partialDataSet, new FileOutputStream(tableToXML.get(tableName)));
			}
		} finally {
			if (connection != null)
				connection.close();
		}

	}

	/**
	 * 增加数据
	 * 增加数据库外键约束
	 *
	 * @param dataSetXML xml格式的数据
	 * @throws Exception 异常
	 */
	public void insertData(File dataSetXML) throws Exception {
		FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
		IDataSet dataSet = builder.build(dataSetXML);
		dataSet = new FilteredDataSet(createFilter(), dataSet);
		DatabaseOperation.INSERT.execute(createConnect(), dataSet);
	}

	/**
	 * 刷新数据
	 * 增加数据库外键约束
	 *
	 * @param dataSetXML
	 * @throws Exception
	 */
	public void refreshData(File dataSetXML) throws Exception {
		FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
		IDataSet dataSet = builder.build(dataSetXML);
		dataSet = new FilteredDataSet(createFilter(), dataSet);
		DatabaseOperation.REFRESH.execute(createConnect(), dataSet);
	}

	/**
	 * 删除数据
	 * 增加数据库外键约束
	 *
	 * @param dataSetXML
	 * @throws Exception
	 */
	public void removeData(File dataSetXML) throws Exception {
		FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
		IDataSet dataSet = builder.build(dataSetXML);
		dataSet = new FilteredDataSet(createFilter(), dataSet);
		DatabaseOperation.DELETE.execute(createConnect(), dataSet);
	}

	public void clearTable(File dataSetXML) throws Exception {
		execute(DatabaseOperation.DELETE_ALL, dataSetXML);
	}

	public void appendData(File dataSetXML) throws Exception {
		execute(DatabaseOperation.INSERT, dataSetXML);
	}

	public void execute(DatabaseOperation operation, File dataSetXML) throws Exception {
		IDatabaseConnection conn = null;
		conn = createConnect();
		FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
		IDataSet dataSet = builder.build(dataSetXML);
		operation.execute(conn, dataSet);
	}


	public String getDriverName() {
		return driverName;
	}

	public void setDriverName(String driverName) {
		this.driverName = driverName;
	}

	public String getJdbcUrl() {
		return jdbcUrl;
	}

	public void setJdbcUrl(String jdbcUrl) {
		this.jdbcUrl = jdbcUrl;
	}

	public String getUsername() {
		return username;
	}

	public void setUsername(String username) {
		this.username = username;
	}

	public String getPassword() {
		return password;
	}

	public void setPassword(String password) {
		this.password = password;
	}

	public Connection getConnection() {
		return connection;
	}

	public void setConnection(Connection connection) {
		this.connection = connection;
	}
}
